import os
import sys
import time
import psutil
import subprocess
import logging
import os
import argparse

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def check_mysql_mcp_server():
    """检查MySQL MCP Server是否正在运行"""
    logger.info("正在检查MySQL MCP Server进程...")
    
    # 获取当前进程ID，排除自身
    current_pid = os.getpid()
    
    # 查找运行server.py的Python进程
    for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
        try:
            if proc.name() == 'python.exe' and proc.pid != current_pid:
                cmdline = ' '.join(proc.cmdline())
                if 'server.py' in cmdline.lower() and 'check_mysql_mcp_server.py' not in cmdline.lower():
                    logger.info(f"发现MySQL MCP Server进程: PID={proc.pid}, 命令行={cmdline}")
                    return True, proc.pid
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass
    
    logger.warning("未找到MySQL MCP Server进程。")
    return False, None

def stop_mysql_mcp_server(pid):
    """停止MySQL MCP Server进程"""
    logger.info(f"尝试停止MySQL MCP Server进程: PID={pid}")
    try:
        proc = psutil.Process(pid)
        # 尝试优雅终止进程
        proc.terminate()
        # 等待进程终止，最多等待5秒
        try:
            proc.wait(timeout=5)
            logger.info(f"MySQL MCP Server进程: PID={pid} 已成功终止。")
            return True
        except psutil.TimeoutExpired:
            # 如果超时，则强制终止进程
            proc.kill()
            logger.warning(f"MySQL MCP Server进程: PID={pid} 超时未终止，已强制终止。")
            return True
    except psutil.NoSuchProcess:
        logger.warning(f"未找到进程: PID={pid}。")
        return False
    except psutil.AccessDenied:
        logger.error(f"无权限终止进程: PID={pid}。")
        return False
    except Exception as e:
        logger.error(f"终止进程: PID={pid} 时出错: {e}")
        return False

def start_mysql_mcp_server():
    """启动MySQL MCP Server"""
    logger.info("尝试启动MySQL MCP Server服务...")
    
    try:
        # 检查当前目录是否存在start_mysql_mcp.bat
        bat_file = 'start_mysql_mcp.bat'
        if not os.path.exists(bat_file):
            logger.error(f"未找到{bat_file}文件，请确保该文件存在于当前目录。")
            return False
        
        # 启动批处理文件，使用新窗口
        subprocess.Popen([bat_file], shell=True, creationflags=subprocess.CREATE_NEW_CONSOLE)
        logger.info(f"已尝试运行{bat_file}启动MySQL MCP Server服务。")
        return True
    except Exception as e:
        logger.error(f"启动MySQL MCP Server服务时出错: {e}")
        return False

if __name__ == '__main__':
    # 检查是否安装了psutil库
    try:
        import psutil
    except ImportError:
        logger.error("未找到psutil库，请先安装: pip install psutil")
        sys.exit(1)
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='检查、启动或重启MySQL MCP Server服务')
    parser.add_argument('--restart', action='store_true', help='重启MySQL MCP Server服务')
    args = parser.parse_args()
    
    # 检查MySQL MCP Server状态
    is_running, pid = check_mysql_mcp_server()
    
    if args.restart:
        logger.info("执行重启MySQL MCP Server服务操作...")
        if is_running:
            # 停止服务
            if stop_mysql_mcp_server(pid):
                # 等待1秒后启动服务
                time.sleep(1)
                if start_mysql_mcp_server():
                    # 等待2秒后再次检查
                    time.sleep(2)
                    is_running, new_pid = check_mysql_mcp_server()
                    if is_running:
                        logger.info("MySQL MCP Server服务已成功重启。")
                    else:
                        logger.warning("MySQL MCP Server服务重启后仍未检测到进程，请手动检查。")
                else:
                    logger.error("重启时无法启动MySQL MCP Server服务。")
            else:
                logger.error("重启时无法停止MySQL MCP Server服务。")
        else:
            logger.info("MySQL MCP Server服务未运行，直接启动...")
            if start_mysql_mcp_server():
                # 等待2秒后再次检查
                time.sleep(2)
                is_running, new_pid = check_mysql_mcp_server()
                if is_running:
                    logger.info("MySQL MCP Server服务已成功启动。")
                else:
                    logger.warning("MySQL MCP Server服务启动后仍未检测到进程，请手动检查。")
            else:
                logger.error("无法启动MySQL MCP Server服务。")
    else:
        if not is_running:
            # 尝试启动服务
            if start_mysql_mcp_server():
                # 等待2秒后再次检查
                time.sleep(2)
                is_running, pid = check_mysql_mcp_server()
                if is_running:
                    logger.info("MySQL MCP Server服务已成功启动。")
                else:
                    logger.warning("MySQL MCP Server服务启动后仍未检测到进程，请手动检查。")
            else:
                logger.error("无法启动MySQL MCP Server服务。")
        else:
            logger.info("MySQL MCP Server服务正在正常运行。")
    
    print("操作完成。")
    input("按Enter键退出...")